{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ab559993-d548-43a8-854a-3892432d8805",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pywt\n",
    "import numpy\n",
    "def format_array(a):\n",
    "    \"\"\"Consistent array representation across different systems\"\"\"\n",
    "    a = numpy.where(numpy.abs(a) < 1e-5, 0, a)\n",
    "    return numpy.array2string(a, precision=5, separator=' ', suppress_small=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d9f94629-572c-41d4-9730-cb734eb05462",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1, 2, 3, 4, 5, 6, 7, 8]\n",
      "''\n",
      "0\n",
      "3\n"
     ]
    }
   ],
   "source": [
    "# Create Wavelet Packet structure\n",
    "x = [1, 2, 3, 4, 5, 6, 7, 8]\n",
    "wp = pywt.WaveletPacket(data=x, wavelet='db1', mode='symmetric')\n",
    "print(wp.data)\n",
    "print(repr(wp.path))\n",
    "print(wp.level)\n",
    "print(wp['ad'].maxlevel)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9f1afc63-2437-4215-b34c-238eff4dfed8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 2.12132034  4.94974747  7.77817459 10.60660172]\n",
      "a\n",
      "[ 5. 13.]\n",
      "aa\n",
      "[12.72792206]\n",
      "aaa\n"
     ]
    }
   ],
   "source": [
    "# Traversing WP tree\n",
    "# 1st level:\n",
    "print(wp['a'].data)\n",
    "print(wp['a'].path)\n",
    "# 2nd level:\n",
    "print(wp['aa'].data)\n",
    "print(wp['aa'].path)\n",
    "# 3rd level:\n",
    "print(wp['aaa'].data)\n",
    "print(wp['aaa'].path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "07404705-55d8-4572-a5a6-8c7ee8a63bd9",
   "metadata": {},
   "outputs": [
    {
     "ename": "IndexError",
     "evalue": "Path length is out of range.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_47046/709472788.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;31m# Ups, we have reached the maximum level of decomposition and got an IndexError:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mwp\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'aaaa'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/pywt/_wavelet_packets.py\u001b[0m in \u001b[0;36m__getitem__\u001b[0;34m(self, path)\u001b[0m\n\u001b[1;32m    262\u001b[0m             if (self.maxlevel is not None and\n\u001b[1;32m    263\u001b[0m                     len(path) > self.maxlevel * self.PART_LEN):\n\u001b[0;32m--> 264\u001b[0;31m                 \u001b[0;32mraise\u001b[0m \u001b[0mIndexError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"Path length is out of range.\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    265\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mpath\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    266\u001b[0m                 return self.get_subnode(path[0:self.PART_LEN], True)[\n",
      "\u001b[0;31mIndexError\u001b[0m: Path length is out of range."
     ]
    }
   ],
   "source": [
    "# Ups, we have reached the maximum level of decomposition and got an IndexError:\n",
    "print(wp['aaaa'].data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "80b7fae1-ed2a-4c7e-a4ce-951eb1871379",
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "Subnode name must be in ['a', 'd'], not 'c'.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_47046/2219559839.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;31m# Now try some invalid path:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mwp\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'ac'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/pywt/_wavelet_packets.py\u001b[0m in \u001b[0;36m__getitem__\u001b[0;34m(self, path)\u001b[0m\n\u001b[1;32m    264\u001b[0m                 \u001b[0;32mraise\u001b[0m \u001b[0mIndexError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"Path length is out of range.\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    265\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mpath\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 266\u001b[0;31m                 return self.get_subnode(path[0:self.PART_LEN], True)[\n\u001b[0m\u001b[1;32m    267\u001b[0m                     path[self.PART_LEN:]]\n\u001b[1;32m    268\u001b[0m             \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/pywt/_wavelet_packets.py\u001b[0m in \u001b[0;36m__getitem__\u001b[0;34m(self, path)\u001b[0m\n\u001b[1;32m    264\u001b[0m                 \u001b[0;32mraise\u001b[0m \u001b[0mIndexError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"Path length is out of range.\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    265\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mpath\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 266\u001b[0;31m                 return self.get_subnode(path[0:self.PART_LEN], True)[\n\u001b[0m\u001b[1;32m    267\u001b[0m                     path[self.PART_LEN:]]\n\u001b[1;32m    268\u001b[0m             \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/pywt/_wavelet_packets.py\u001b[0m in \u001b[0;36mget_subnode\u001b[0;34m(self, part, decompose)\u001b[0m\n\u001b[1;32m    226\u001b[0m             \u001b[0;34m(\u001b[0m\u001b[0mdefault\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    227\u001b[0m         \"\"\"\n\u001b[0;32m--> 228\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_validate_node_name\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpart\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    229\u001b[0m         \u001b[0msubnode\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_get_node\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpart\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    230\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0msubnode\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mdecompose\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mis_empty\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/pywt/_wavelet_packets.py\u001b[0m in \u001b[0;36m_validate_node_name\u001b[0;34m(self, part)\u001b[0m\n\u001b[1;32m    110\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_validate_node_name\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpart\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    111\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mpart\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mPARTS\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 112\u001b[0;31m             raise ValueError(\"Subnode name must be in [%s], not '%s'.\" %\n\u001b[0m\u001b[1;32m    113\u001b[0m                              (', '.join(\"'%s'\" % p for p in self.PARTS), part))\n\u001b[1;32m    114\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mValueError\u001b[0m: Subnode name must be in ['a', 'd'], not 'c'."
     ]
    }
   ],
   "source": [
    "# Now try some invalid path:\n",
    "print(wp['ac'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d0da0ee3-d48a-4309-a655-5a66f2b873e4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-2. -2.]\n",
      "ad\n",
      "d\n",
      "a\n",
      "2\n",
      "3\n",
      "symmetric\n"
     ]
    }
   ],
   "source": [
    "# Accessing Node’s attributes:\n",
    "print(wp['ad'].data)\n",
    "print(wp['ad'].path)\n",
    "print(wp['ad'].node_name)\n",
    "print(wp['ad'].parent.path)\n",
    "print(wp['ad'].level)\n",
    "print(wp['ad'].maxlevel)\n",
    "print(wp['ad'].mode)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6c233e52-28f5-4e01-ab46-cbb6bfcd37a5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['aaa', 'aad', 'ada', 'add', 'daa', 'dad', 'dda', 'ddd']\n"
     ]
    }
   ],
   "source": [
    "# Collecting nodes\n",
    "# We can get all nodes on the particular level either in natural order:\n",
    "print([node.path for node in wp.get_level(3, 'natural')])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c14eb482-719d-43c4-8fcc-3f53e1774eb9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['aaa', 'aad', 'add', 'ada', 'dda', 'ddd', 'dad', 'daa']\n"
     ]
    }
   ],
   "source": [
    "# or sorted based on the band frequency (freq):\n",
    "print([node.path for node in wp.get_level(3, 'freq')])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "66acc45a-06c2-4550-98e7-dccb0f76e245",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1. 2. 3. 4. 5. 6. 7. 8.]\n",
      "None\n",
      "[1. 2. 3. 4. 5. 6. 7. 8.]\n",
      "[1. 2. 3. 4. 5. 6. 7. 8.]\n",
      "['aa', 'ad', 'd']\n",
      "['aaa', 'aad', 'ada', 'add', 'daa', 'dad', 'dda', 'ddd']\n"
     ]
    }
   ],
   "source": [
    "# Reconstructing data from Wavelet Packets:\n",
    "# Now create a new Wavelet Packet and set its nodes with some data.\n",
    "new_wp = pywt.WaveletPacket(data=None, wavelet='db1', mode='symmetric')\n",
    "new_wp['aa'] = wp['aa'].data\n",
    "new_wp['ad'] = [-2., -2.]\n",
    "# For convenience, Node.data gets automatically extracted from the Node object:\n",
    "new_wp['d'] = wp['d']\n",
    "# And reconstruct the data from the aa, ad and d packets.\n",
    "print(new_wp.reconstruct(update=False))\n",
    "# If the update param in the reconstruct method is set to False, the node’s data will not be updated.\n",
    "print(new_wp.data)\n",
    "# Otherwise, the data attribute will be set to the reconstructed value.\n",
    "print(new_wp.reconstruct(update=True))\n",
    "print(new_wp.data)\n",
    "print([n.path for n in new_wp.get_leaf_nodes(False)])\n",
    "print([n.path for n in new_wp.get_leaf_nodes(True)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "40728280-4ebe-402c-8ff4-3eda3b484de2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "aaa [12.72792]\n",
      "aad [-5.65685]\n",
      "ada [-2.82843]\n",
      "add [0.]\n",
      "daa [-1.41421]\n",
      "dad [0.]\n",
      "dda [0.]\n",
      "ddd [0.]\n"
     ]
    }
   ],
   "source": [
    "# Removing nodes from Wavelet Packet tree\n",
    "dummy = wp.get_level(2)\n",
    "for n in wp.get_leaf_nodes(False):\n",
    "    print(n.path, format_array(n.data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "54af5a28-e9d2-499e-aa32-142f158eae40",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ad: [-2. -2.]\n"
     ]
    }
   ],
   "source": [
    "node = wp['ad']\n",
    "print(node)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "c0b43ee3-7ec5-4cf6-be0c-b68e874b6bdf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "aaa [12.72792]\n",
      "aad [-5.65685]\n",
      "daa [-1.41421]\n",
      "dad [0.]\n",
      "dda [0.]\n",
      "ddd [0.]\n"
     ]
    }
   ],
   "source": [
    "# To remove a node from the WP tree, use Python’s del obj[x] (Node.__delitem__):\n",
    "del wp['ad']\n",
    "for n in wp.get_leaf_nodes():\n",
    "    print(n.path, format_array(n.data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "950ca355-e414-40d9-ad95-b062f254ab87",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2. 3. 2. 3. 6. 7. 6. 7.]\n"
     ]
    }
   ],
   "source": [
    "print(wp.reconstruct())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "7c68f61d-bccf-40b0-95a2-be61be2e28ab",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "aa [ 5. 13.]\n",
      "ad [-2. -2.]\n",
      "daa [-1.41421]\n",
      "dad [0.]\n",
      "dda [0.]\n",
      "ddd [0.]\n"
     ]
    }
   ],
   "source": [
    "# Now restore the deleted node value.\n",
    "wp['ad'].data = node.data\n",
    "for n in wp.get_leaf_nodes(False):\n",
    "    print(n.path, format_array(n.data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "ccc25075-a462-490c-9098-fab265f10725",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1. 2. 3. 4. 5. 6. 7. 8.]\n"
     ]
    }
   ],
   "source": [
    "print(wp.reconstruct())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "29e38c87-8975-4d36-bc44-29ca3a40388b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Lazy evaluation:\n",
    "x = [1, 2, 3, 4, 5, 6, 7, 8]\n",
    "wp = pywt.WaveletPacket(data=x, wavelet='db1', mode='symmetric')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "34cd5ef2-f13d-4175-8819-15b64625a51a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "None\n",
      "None\n"
     ]
    }
   ],
   "source": [
    "# 1. At first the wp’s attribute a is None\n",
    "print(wp.a)\n",
    "print(wp.d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "f5900e5a-2be8-4eba-b7e7-487214c015e1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a: [ 2.12132034  4.94974747  7.77817459 10.60660172]\n"
     ]
    }
   ],
   "source": [
    "# 2. At first attempt to access the node it is computed via decomposition of its parent node (the wp object itself).\n",
    "print(wp['a'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "6e9832e7-2dad-4e61-be41-e64ea8006e64",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a: [ 2.12132034  4.94974747  7.77817459 10.60660172]\n",
      "d: [-0.70710678 -0.70710678 -0.70710678 -0.70710678]\n"
     ]
    }
   ],
   "source": [
    "# 3. Now the wp.a is set to the newly created node:\n",
    "print(wp.a)\n",
    "# And so is wp.d:\n",
    "print(wp.d)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
